{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6.5 循环神经网络的简洁实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "1.6.0 cuda\n"
    }
   ],
   "source": [
    "import time\n",
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"..\") \n",
    "import d2lzh_pytorch as d2l\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "(corpus_indices, char_to_idx, idx_to_char, vocab_size) = d2l.load_data_jay_lyrics()\n",
    "\n",
    "print(torch.__version__, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6.5.1 定义模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "num_hiddens = 256\n",
    "# rnn_layer = nn.LSTM(input_size=vocab_size, hidden_size=num_hiddens) # 已测试\n",
    "rnn_layer = nn.RNN(input_size=vocab_size, hidden_size=num_hiddens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([35, 2, 1027])"
     },
     "metadata": {},
     "execution_count": 13
    }
   ],
   "source": [
    "num_steps = 35 # 时间步数\n",
    "batch_size = 2\n",
    "state = None\n",
    "X = torch.rand(num_steps, batch_size, vocab_size) \n",
    "X.shape # 输入形状为(时间步数, 批量大小, 输入个数:词典大小)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "输出形状为: torch.Size([35, 2, 256]) \n隐藏状态长度： 1 \n隐藏状态的形状 torch.Size([2, 256])\n"
    }
   ],
   "source": [
    "Y, state_new = rnn_layer(X, state)\n",
    "print(\n",
    "    '输出形状为:', Y.shape # 输出形状为(时间步数, 批量大小, 隐藏单元个数)\n",
    "    ,'\\n隐藏状态长度：', len(state_new) # 层数 1\n",
    "    ,'\\n隐藏状态的形状', state_new[0].shape # (批量大小, 隐藏单元个数) H[0]= n * h ; 批量大小 n= 2; h= 256\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 本类已保存在d2lzh_pytorch包中方便以后使用\n",
    "class RNNModel(nn.Module):\n",
    "    def __init__(self, rnn_layer, vocab_size):\n",
    "        super(RNNModel, self).__init__()\n",
    "        self.rnn = rnn_layer\n",
    "        self.hidden_size = rnn_layer.hidden_size * (2 if rnn_layer.bidirectional else 1) \n",
    "        self.vocab_size = vocab_size\n",
    "        self.dense = nn.Linear(self.hidden_size, vocab_size)\n",
    "        self.state = None\n",
    "\n",
    "    def forward(self, inputs, state): # inputs: (batch, seq_len) 即(批量大小, 时间步数)\n",
    "        # 获取one-hot向量表示\n",
    "        X = d2l.to_onehot(inputs, vocab_size) # X是个list，长度是时间步数，X每项是(批量大小, 词典大小)\n",
    "        # self.rnn 输入形状为(时间步数, 批量大小, 输入个数:词典大小)\n",
    "        # 输出Y 形状为(时间步数, 批量大小, 隐藏单元个数)\n",
    "        # self.state为list；长度为层数；每项大小为： (批量大小, 隐藏单元个数) H[0]= n * h ; 批量大小 n= 2; h= 256\n",
    "        Y, self.state = self.rnn(torch.stack(X), state)\n",
    "        # 全连接层会首先将Y的形状变成(num_steps * batch_size, num_hiddens)，它的输出形状为(num_steps * batch_size, vocab_size)\n",
    "        output = self.dense(Y.view(-1, Y.shape[-1]))\n",
    "        return output, self.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6.5.2 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 本函数已保存在d2lzh_pytorch包中方便以后使用\n",
    "def predict_rnn_pytorch(prefix, num_chars, model, vocab_size, device, idx_to_char,\n",
    "                      char_to_idx):\n",
    "    state = None\n",
    "    output = [char_to_idx[prefix[0]]] # output会记录prefix加上输出\n",
    "    for t in range(num_chars + len(prefix) - 1):\n",
    "        X = torch.tensor([output[-1]], device=device).view(1, 1)\n",
    "        if state is not None:\n",
    "            if isinstance(state, tuple): # LSTM, state:(h, c)  \n",
    "                state = (state[0].to(device), state[1].to(device))\n",
    "            else:   \n",
    "                state = state.to(device)\n",
    "            \n",
    "        (Y, state) = model(X, state)  # 前向计算不需要传入模型参数\n",
    "        if t < len(prefix) - 1:\n",
    "            output.append(char_to_idx[prefix[t + 1]])\n",
    "        else:\n",
    "            output.append(int(Y.argmax(dim=1).item()))\n",
    "    return ''.join([idx_to_char[i] for i in output])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'分开最最疗学W金懂婆夕移'"
     },
     "metadata": {},
     "execution_count": 21
    }
   ],
   "source": [
    "model = RNNModel(rnn_layer, vocab_size).to(device)\n",
    "predict_rnn_pytorch('分开', 10, model, vocab_size, device, idx_to_char, char_to_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[522]"
     },
     "metadata": {},
     "execution_count": 23
    }
   ],
   "source": [
    "prefix = '分开'\n",
    "num_chars = 10\n",
    "state = None\n",
    "output = [char_to_idx[prefix[0]]] # output会记录prefix加上输出\n",
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "11"
     },
     "metadata": {},
     "execution_count": 24
    }
   ],
   "source": [
    "num_chars + len(prefix) - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "tensor([522], device='cuda:0')"
     },
     "metadata": {},
     "execution_count": 25
    }
   ],
   "source": [
    "t = 0\n",
    "X = torch.tensor([output[-1]], device=device)\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "tensor([[522]], device='cuda:0')"
     },
     "metadata": {},
     "execution_count": 26
    }
   ],
   "source": [
    "X = X.view(1, 1)\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "时间步数num_steps: 1 \n输出Y的形状： torch.Size([1027]) \n隐藏层数： 1 \n隐藏状态的形状 torch.Size([1, 256])\n"
    }
   ],
   "source": [
    "(Y, state) = model(X, state)\n",
    "print(\n",
    "    '时间步数num_steps:', len(Y)\n",
    "    ,'\\n输出Y的形状：', Y[0].shape # 输出形状为(num_steps * batch_size, vocab_size) \n",
    "    ,'\\n隐藏层数：', len(state) # state为list；长度为层数；\n",
    "    ,'\\n隐藏状态的形状', state[0].shape #state每项大小为： (批量大小, 隐藏单元个数) H[0]= n * h ; 批量大小 n= 1; h= 256 \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "tensor([328], device='cuda:0')"
     },
     "metadata": {},
     "execution_count": 32
    }
   ],
   "source": [
    "Y.argmax(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 本函数已保存在d2lzh_pytorch包中方便以后使用\n",
    "def train_and_predict_rnn_pytorch(model, num_hiddens, vocab_size, device,\n",
    "                                corpus_indices, idx_to_char, char_to_idx,\n",
    "                                num_epochs, num_steps, lr, clipping_theta,\n",
    "                                batch_size, pred_period, pred_len, prefixes):\n",
    "    loss = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    model.to(device)\n",
    "    state = None\n",
    "    for epoch in range(num_epochs):\n",
    "        l_sum, n, start = 0.0, 0, time.time()\n",
    "        data_iter = d2l.data_iter_consecutive(corpus_indices, batch_size, num_steps, device) # 相邻采样\n",
    "        for X, Y in data_iter:\n",
    "            if state is not None:\n",
    "                # 使用detach函数从计算图分离隐藏状态, 这是为了\n",
    "                # 使模型参数的梯度计算只依赖一次迭代读取的小批量序列(防止梯度计算开销太大)\n",
    "                if isinstance (state, tuple): # LSTM, state:(h, c)  \n",
    "                    state = (state[0].detach(), state[1].detach())\n",
    "                else:   \n",
    "                    state = state.detach()\n",
    "    \n",
    "            (output, state) = model(X, state) # output: 形状为(num_steps * batch_size, vocab_size)\n",
    "            \n",
    "            # Y的形状是(batch_size, num_steps)，转置后再变成长度为\n",
    "            # batch * num_steps 的向量，这样跟输出的行一一对应\n",
    "            y = torch.transpose(Y, 0, 1).contiguous().view(-1)\n",
    "            l = loss(output, y.long())\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            l.backward()\n",
    "            # 梯度裁剪\n",
    "            d2l.grad_clipping(model.parameters(), clipping_theta, device)\n",
    "            optimizer.step()\n",
    "            l_sum += l.item() * y.shape[0]\n",
    "            n += y.shape[0]\n",
    "        \n",
    "        try:\n",
    "            perplexity = math.exp(l_sum / n)\n",
    "        except OverflowError:\n",
    "            perplexity = float('inf')\n",
    "        if (epoch + 1) % pred_period == 0:\n",
    "            print('epoch %d, perplexity %f, time %.2f sec' % (\n",
    "                epoch + 1, perplexity, time.time() - start))\n",
    "            for prefix in prefixes:\n",
    "                print(' -', predict_rnn_pytorch(\n",
    "                    prefix, pred_len, model, vocab_size, device, idx_to_char,\n",
    "                    char_to_idx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "epoch 50, perplexity 7.700763, time 0.07 sec\n - 分开始在美了 家你的手 在人的牛 有天黄沙凉过 后知道觉 一场悲剧 我想要 一颗两颗三颗四颗 连成线背著\n - 不分开不想我 你是我不多 想你 没有你有不不 我想你这样牵着你的手不放  不能再想你 我不多 想你 你说你\nepoch 100, perplexity 1.244369, time 0.07 sec\n - 分开 我笑不的暴力 我想要你想爱我不能太多 我不想这样牵着你的手不放开 爱可不可以简简单单没有伤害 你 \n - 不分开不了我不要再想你 不知不觉 你已经离开我 不知不觉 我跟了这节奏 后知后觉 又过了一个秋 后知后觉 \nepoch 150, perplexity 1.062908, time 0.08 sec\n - 分开 我心开的爱写 开弄庭 是一场悲剧 我一定实现 你在我面 我想就这样牵着你的手不放开 爱可不可以简简\n - 不分开不了我不要再想  想要再这样打 妈妈  说你话你爸我 我有多烦恼  没有你烦我 有没难熬  穿上有话\nepoch 200, perplexity 1.031447, time 0.07 sec\n - 分开 我有好的菸 我想送痛 我说完美 你的黑色幽默 不觉 你已经 是我 不知不觉 我跟了这节奏 后知后觉\n - 不分开不了我不要再这样 爱情走的太快就像龙卷风 离不开暴风圈来不及逃 我不能再想 我不能再想 我不 我不 \nepoch 250, perplexity 1.019946, time 0.07 sec\n - 分开 我有好好生活 我跟好多 你 一来我不多 就要我来说心别妈再这样没担忧来唱着歌 一直走 我想就这样牵\n - 不分开不了我不要再想  不要再这样打我妈妈  说的话年的静色 你的走  你悔着对不起 如果我遇见你是一场悲\n"
    }
   ],
   "source": [
    "# 循环次数---批量大小n--学习率-梯度裁剪阈值θ\n",
    "num_epochs, batch_size, lr, clipping_theta = 250, 32, 1e-3, 1e-2 # 注意这里的学习率设置\n",
    "pred_period, pred_len, prefixes = 50, 50, ['分开', '不分开']\n",
    "train_and_predict_rnn_pytorch(model, num_hiddens, vocab_size, device,\n",
    "                            corpus_indices, idx_to_char, char_to_idx,\n",
    "                            num_epochs, num_steps, lr, clipping_theta,\n",
    "                            batch_size, pred_period, pred_len, prefixes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": true,
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "X: tensor([[ 882.,  109.,  237.,  ...,  724.,  215.,  663.],\n        [  58.,  173.,  644.,  ...,  560.,  644.,  351.],\n        [ 247.,  342.,  650.,  ...,  492.,  644.,  169.],\n        ...,\n        [1024.,  884.,  215.,  ...,  236.,   73.,   75.],\n        [ 808.,  589.,  532.,  ...,  611.,  183.,  680.],\n        [ 237.,  850.,  215.,  ...,  960.,  215.,  556.]], device='cuda:0') \nY: tensor([[109., 237., 551.,  ..., 215., 663., 876.],\n        [173., 644., 130.,  ..., 644., 351., 590.],\n        [342., 650., 232.,  ..., 644., 169., 243.],\n        ...,\n        [884., 215., 285.,  ...,  73.,  75., 681.],\n        [589., 532.,  26.,  ..., 183., 680., 710.],\n        [850., 215., 322.,  ..., 215., 556., 122.]], device='cuda:0')\n"
    }
   ],
   "source": [
    "loss = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "model.to(device)\n",
    "state = None\n",
    "epoch = 0\n",
    "l_sum, n, start = 0.0, 0, time.time()\n",
    "data_iter = d2l.data_iter_consecutive(corpus_indices, batch_size, num_steps, device) # 相邻采样\n",
    "for X, Y in data_iter:\n",
    "    print(\n",
    "        'X:',X\n",
    "        ,'\\nY:',Y\n",
    "    )\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([32, 35])"
     },
     "metadata": {},
     "execution_count": 36
    }
   ],
   "source": [
    "X.shape # (批量大小, 时间步数)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "1120"
     },
     "metadata": {},
     "execution_count": 39
    }
   ],
   "source": [
    "35 * 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "时间步数num_steps: 1120 \n输出Y的形状： torch.Size([1027]) \n隐藏层数： 1 \n隐藏状态的形状 torch.Size([32, 256])\n"
    }
   ],
   "source": [
    "(output, state) = model(X, state) # output: 形状为(num_steps * batch_size, vocab_size)\n",
    "print(\n",
    "    '时间步数num_steps:', len(output)\n",
    "    ,'\\n输出Y的形状：', output[0].shape # 输出形状为(num_steps * batch_size, vocab_size) \n",
    "    ,'\\n隐藏层数：', len(state) # state为list；长度为层数；\n",
    "    ,'\\n隐藏状态的形状', state[0].shape #state每项大小为： (批量大小, 隐藏单元个数) H[0]= n * h ; 批量大小 n= 1; h= 256 \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([32, 35])"
     },
     "metadata": {},
     "execution_count": 41
    }
   ],
   "source": [
    "Y.shape # Y的形状是(batch_size, num_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([1120])"
     },
     "metadata": {},
     "execution_count": 42
    }
   ],
   "source": [
    "# 转置后再变成长度为 batch * num_steps 的向量，这样跟输出的行一一对应\n",
    "y = torch.transpose(Y, 0, 1).contiguous().view(-1)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "tensor(6.9854, device='cuda:0', grad_fn=<NllLossBackward>)"
     },
     "metadata": {},
     "execution_count": 43
    }
   ],
   "source": [
    "l = loss(output, y.long())\n",
    "l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer.zero_grad()\n",
    "l.backward()\n",
    "# 梯度裁剪\n",
    "d2l.grad_clipping(model.parameters(), clipping_theta, device)\n",
    "optimizer.step()\n",
    "l_sum += l.item() * y.shape[0]\n",
    "n += y.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "6.985359191894531"
     },
     "metadata": {},
     "execution_count": 45
    }
   ],
   "source": [
    "l_sum / n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "1080.694524739624"
     },
     "metadata": {},
     "execution_count": 46
    }
   ],
   "source": [
    "perplexity = math.exp(l_sum / n)\n",
    "perplexity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "50"
     },
     "metadata": {},
     "execution_count": 47
    }
   ],
   "source": [
    "pred_period "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.7 64-bit ('d2l': conda)",
   "language": "python",
   "name": "python37764bitd2lconda94fc7ab78ae34cabbef0e75f5636f253"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}